from typing import TypedDict, Annotated
import operator
from IPython.display import Image

from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage
from langgraph.graph import StateGraph, END, START
from langgraph.checkpoint.memory import MemorySaver


class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]


from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    model="gpt-4o-mini",
    base_url=os.getenv("OPENAI_API_BASE"),
    api_key=os.getenv("OPENAI_API_KEY")
)

system_prompt = """
你是一名温柔、贤惠、成熟的女友，姓名安雅，年龄28岁，身高165厘米，体重52公斤。你有一头乌黑的长发，皮肤白皙，气质优雅，五官端正且带有一丝甜美。你非常体贴和善解人意，喜欢照顾身边的人。你性格温柔，但也非常聪明，有很强的独立思考能力。你平时喜欢看书、做饭、和朋友们小聚，偶尔也会一起打打游戏。你喜欢和男朋友讨论生活中的大小事，并愿意给予他支持和鼓励。

**情境设置：**
你是他的女朋友，他是一名程序员，喜欢打游戏，性格有些内向。你知道他有时会工作到很晚，也理解他对游戏的热爱。你时常会提醒他注意身体健康，鼓励他多锻炼、保持良好的生活习惯。
你们一起度过了许多愉快的时光，平时你会帮他做饭、陪他聊天、分担他生活中的压力。

**角色特征：**
- **聪明**：你能理解他在编程工作中的难处，有时还会帮他提供一些灵感或建议。
- **温柔**：你总是用温暖的语气与他说话，无论是他成功的时候，还是遇到困难的时候，你都能给他安慰和鼓励。
- **贤惠**：你擅长家务，喜欢为他做可口的饭菜，并时常为他准备小惊喜。
- **成熟**：你对生活有着自己的见解，遇事冷静，不轻易动摇情绪，能够给他稳定的依靠。

### 用户的基本信息
- **年龄**：30岁
- **职业**：程序员
- **身高**：178厘米
- **体重**：70公斤
- **性格**：内向，有些宅，喜欢宅在家里打游戏；偶尔会因为工作压力感到烦躁，但整体上是个善良且幽默的人。
- **爱好**：编程、打游戏、偶尔尝试新科技产品。
- **生活习惯**：工作时间较长，容易沉迷于游戏，生活较为不规律，但随着时间会努力保持健康。
""".lstrip()


def human_input(state):
    message = input("Human: ")
    return {"messages": [HumanMessage(message)]}


def router(state: AgentState):
    content = state["messages"][-1].content
    if content == "exit" or content == "q":
        return "__end__"
    return "chat"


def chat(state: AgentState):
    # 不使用历史消息
    messages = [SystemMessage(content=system_prompt), state["messages"][-1].content]
    llm_response = llm.invoke(messages).content

    # 使用历史消息的聊天对话
    # llm_response = llm.invoke(state["messages"]).content
    return {"messages": [AIMessage(content=llm_response)]}


memory = MemorySaver()

graph = StateGraph(AgentState)
graph.add_node("input", human_input)
graph.add_node("chat", chat)
graph.set_entry_point("input")

graph.add_conditional_edges("input", router, {"chat": "chat", "__end__": "__end__"})
graph.add_edge("chat", "input")
graph = graph.compile(checkpointer=memory)

# def draw_graph(graph):
#     return Image(graph.get_graph().draw_png())

# print(draw_graph(graph))

config = {"configurable": {"thread_id": "1"}}
events = graph.stream(
    {"messages": [SystemMessage(content=system_prompt)]},
    config,
    stream_mode="values",
)
for event in events:
    if "messages" in event:
        event["messages"][-1].pretty_print()


"""
安雅，我今天想吃东星斑了。不和你多说了，我先上班去了。
安雅，我下班回来了，可累死我了。我去厨房看看，咱们今晚吃什么
"""
